import chromadb
import pymysql
from pymysqlreplication import BinLogStreamReader
from pymysqlreplication.row_event import WriteRowsEvent
import time


# MySQL连接信息
def get_connection(host, user, password, db):
    # 创建数据库连接
    conn = pymysql.connect(
        host=host,
        user=user,
        password=password,
        database=db,
        port=3306,
        # charset='utf8'
    )
    return conn


# 一次全量读取
def import_data_from_table():
    connection = get_connection("192.168.0.153","root","Cosmo_fuwu666","text_to_sql")
    cursor = connection.cursor()

    cursor.execute('SELECT question FROM t_audited_info')
    data = [row[0] for row in cursor.fetchall()]

    chroma_client = chromadb.HttpClient("192.168.22.101", 8000)

    collection = chroma_client.get_or_create_collection(name="my_collection")
    collection.add(
        # 待定字段数量
        ids=["id4", "id5", "id6", "id7"],
        metadatas=[{"source": "my_source"}, {"source": "my_source"}, {"source": "my_source"}, {"source": "my_source"}],
        documents=data
    )

    cursor.close()
    connection.close()

    return data


#监控binlog日志，同步到chromadb
def monitor_binlog_for_new_data():
    mysql_settings = {
        "host": "192.168.0.153",
        "port": 3306,
        "user": "root",
        "passwd": "Cosmo_fuwu666",
        # "charset": "utf8"
    }
    stream = BinLogStreamReader(
        connection_settings=mysql_settings,
        blocking=True,
        server_id=1001,
        only_schemas=["text_to_sql"],
        only_tables=["t_audited_info"],
        only_events=[WriteRowsEvent],
    )
    chroma_client = chromadb.HttpClient("192.168.22.101", 8000)
    for binlog_event in stream:
        for row in binlog_event.rows:
            values = row['values']
            print(values)
            # time.sleep(2)
            # 获取新插入的查询
            question = str(values.get("question"))
            idn = str(values.get("id"))

            collection = chroma_client.get_or_create_collection(name="my_collection")
            collection.add(
                ids=[idn],  # 目前这个参数还不确定
                metadatas=[{"source": "my_source"}],  # 目前参数个数不确定
                documents=[question]
            )
    stream.close()
    return None


if __name__ == "__main__":
    # 第一步：全量读取数据
    # data_from_table = import_data_from_table()
    # print("Data from table:", data_from_table)

    # 第二步：监控 binlog
    new_data_from_binlog = monitor_binlog_for_new_data()
    print("New data from binlog:", new_data_from_binlog)
